import unittest
from datetime import datetime, timedelta
from backend.utils.reboot_analyzer import analyze_reboot_times_and_assign_stages

class TestRebootAnalyzer(unittest.TestCase):

    def create_log_entry(self, dt_str: str, message: str, full_message: str = None) -> dict:
        """Helper to create a log entry dict."""
        return {
            'datetime': dt_str,
            'message': message,
            'full_message': full_message if full_message else f"{dt_str} {message}"
        }

    def test_no_reboot(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.000", "Normal log 1"),
            self.create_log_entry("2023-01-01 10:00:01.000", "Normal log 2"),
            self.create_log_entry("2023-01-01 10:00:02.000", "Normal log 3"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 0)
        self.assertEqual(len(result['fixed_logs']), 3)
        self.assertTrue(all(log['run_stage_id'] == 0 for log in result['fixed_logs']))

    def test_1970_reboot_simple(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.000", "Normal log before reboot"),
            self.create_log_entry("1970-01-01 00:00:01.000", "Log at 1970 (reboot start)"),
            self.create_log_entry("1970-01-01 00:00:02.000", "Log in 1970"),
            self.create_log_entry("2023-01-01 11:00:00.000", "Normal log after reboot"),
            self.create_log_entry("2023-01-01 11:00:01.000", "Normal log after reboot 2"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 1)
        self.assertEqual(len(result['fixed_logs']), 5)
        # 验证 run_stage_id
        self.assertEqual(result['fixed_logs'][0]['run_stage_id'], 0)
        self.assertEqual(result['fixed_logs'][1]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][2]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][3]['run_stage_id'], 1) # 1970 到 2023 的跳变属于同一阶段修复
        self.assertEqual(result['fixed_logs'][4]['run_stage_id'], 1)
        
        # 验证修复后的时间
        self.assertTrue(result['fixed_logs'][1]['is_fixed_time'])
        self.assertTrue(datetime.strptime(result['fixed_logs'][1]['datetime'], '%Y-%m-%d %H:%M:%S.%f').year >= 2000)

    def test_1970_reboot_with_internal_rollback(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.000", "Normal log before reboot"),
            self.create_log_entry("1970-01-01 00:00:05.000", "Log 1 in 1970"),
            self.create_log_entry("1970-01-01 00:00:03.000", "Log 2 in 1970 (rollback, new stage)"),
            self.create_log_entry("1970-01-01 00:00:04.000", "Log 3 in 1970"),
            self.create_log_entry("2023-01-01 11:00:00.000", "Normal log after reboot"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 2) # 内部回滚会触发新的重启事件
        self.assertEqual(len(result['fixed_logs']), 5)
        self.assertEqual(result['fixed_logs'][0]['run_stage_id'], 0)
        self.assertEqual(result['fixed_logs'][1]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][2]['run_stage_id'], 2) # 内部回滚后是新的阶段
        self.assertEqual(result['fixed_logs'][3]['run_stage_id'], 2)
        self.assertEqual(result['fixed_logs'][4]['run_stage_id'], 2)

    def test_specific_log_reboot(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.000", "Normal log"),
            self.create_log_entry("2023-01-01 10:00:01.000", "[:00000002] framework: LAUNCH snlua bootstrap (reboot trigger)"),
            self.create_log_entry("2023-01-01 10:00:02.000", "Log after trigger"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 1)
        self.assertEqual(len(result['fixed_logs']), 3)
        self.assertEqual(result['fixed_logs'][0]['run_stage_id'], 0)
        self.assertEqual(result['fixed_logs'][1]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][2]['run_stage_id'], 1)

    def test_multiple_reboots(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.000", "Normal log stage 0"),
            self.create_log_entry("1970-01-01 00:00:01.000", "Reboot 1 start (1970)"),
            self.create_log_entry("2023-01-01 11:00:00.000", "Reboot 1 end (normal time)"),
            self.create_log_entry("2023-01-01 11:00:01.000", "Normal log stage 1"),
            self.create_log_entry("2023-01-01 12:00:00.000", "[:00000002] framework: LAUNCH snlua bootstrap (Reboot 2 trigger)"),
            self.create_log_entry("2023-01-01 12:00:01.000", "Normal log after Reboot 2"),
            self.create_log_entry("1970-01-01 00:00:05.000", "Reboot 3 start (1970)"),
            self.create_log_entry("1970-01-01 00:00:03.000", "Reboot 3 internal rollback"),
            self.create_log_entry("2023-01-01 13:00:00.000", "Reboot 3 end (normal time)"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 4) # 第一次1970、第二次特定日志、第三次1970内部回滚、第四次1970到正常
        self.assertEqual(len(result['fixed_logs']), 9)
        # 验证 run_stage_id
        self.assertEqual(result['fixed_logs'][0]['run_stage_id'], 0) # Stage 0
        self.assertEqual(result['fixed_logs'][1]['run_stage_id'], 1) # Stage 1 (1970 reboot)
        self.assertEqual(result['fixed_logs'][2]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][3]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][4]['run_stage_id'], 2) # Stage 2 (specific log reboot)
        self.assertEqual(result['fixed_logs'][5]['run_stage_id'], 2)
        self.assertEqual(result['fixed_logs'][6]['run_stage_id'], 3) # Stage 3 (1970 reboot)
        self.assertEqual(result['fixed_logs'][7]['run_stage_id'], 4) # Stage 4 (1970 internal rollback)
        self.assertEqual(result['fixed_logs'][8]['run_stage_id'], 4)

    def test_mixed_timestamp_formats(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.123", "Log with microseconds"),
            self.create_log_entry("2023-01-01 10:00:01", "Log without microseconds"),
            self.create_log_entry("1970-01-01 00:00:00.000", "1970 log with microseconds"),
            self.create_log_entry("1970-01-01 00:00:01", "1970 log without microseconds"),
            self.create_log_entry("2023-01-01 11:00:00.000", "Normal log after 1970"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 1)
        self.assertEqual(len(result['fixed_logs']), 5)
        self.assertEqual(result['fixed_logs'][0]['run_stage_id'], 0)
        self.assertEqual(result['fixed_logs'][1]['run_stage_id'], 0)
        self.assertEqual(result['fixed_logs'][2]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][3]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][4]['run_stage_id'], 1)
        # 验证修复后的时间
        self.assertTrue(datetime.strptime(result['fixed_logs'][2]['datetime'], '%Y-%m-%d %H:%M:%S.%f').year >= 2000)

    def test_logs_at_end_of_1970_period(self):
        logs = [
            self.create_log_entry("2023-01-01 10:00:00.000", "Normal log"),
            self.create_log_entry("1970-01-01 00:00:01.000", "Log in 1970 period"),
            self.create_log_entry("1970-01-01 00:00:02.000", "Last log in 1970 period"),
        ]
        result = analyze_reboot_times_and_assign_stages(logs)
        self.assertEqual(len(result['reboot_events']), 1) # 即使没有跳回正常时间，也应视为一次重启
        self.assertEqual(len(result['fixed_logs']), 3)
        self.assertEqual(result['fixed_logs'][0]['run_stage_id'], 0)
        self.assertEqual(result['fixed_logs'][1]['run_stage_id'], 1)
        self.assertEqual(result['fixed_logs'][2]['run_stage_id'], 1) 